import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.examples.tutorials.mnist import input_data


mnist = input_data.read_data_sets(
    'MNIST_data/', validation_size=0, one_hot=True)

print('train_images_type:', type(mnist.train.images))
print('train_images_shape:', mnist.train.images.shape)
print('train_labels_type:', type(mnist.train.labels))
print('train_labels_shape:', mnist.train.labels.shape)
print('test_images_type:', type(mnist.test.images))
print('test_images_shape:', mnist.test.images.shape)
print('test_labels_type:', type(mnist.test.labels))
print('test_labels_shape:', mnist.test.labels.shape)

image = mnist.train.images[0]
image = image.reshape(28, 28)

label = mnist.train.labels[0]

plt.imshow(image, cmap='Greys_r')
plt.show()
print(label)

inputs = tf.placeholder(tf.float32, [None, 784])

labels = tf.placeholder(tf.float32, [None, 10])

fc_1 = tf.layers.dense(inputs, 100, activation=tf.nn.relu)

predict = tf.layers.dense(fc_1, 10)

loss = tf.reduce_mean(tf.square(predict-labels))

optimizer = tf.\
    train.GradientDescentOptimizer(0.1).minimize(loss)

correct_predicts = tf.\
    equal(tf.argmax(predict, 1), tf.argmax(labels, 1))

accuracy = tf\
    .reduce_mean(tf.cast(correct_predicts, tf.float32))

batch_size = 200
epochs = 20
n_batches = mnist.train.num_examples // batch_size
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
print('Initialized')

for epoch in range(epochs):
    for batch in range(n_batches):
        batch_x, batch_y = mnist.train.next_batch(batch_size)
        train_dict = {inputs: batch_x, labels: batch_y}
        sess.run(optimizer, feed_dict=train_dict)
        
    val_feed = {inputs: mnist.test.images,
                labels: mnist.test.labels}
    acc = sess.run(accuracy, feed_dict=val_feed)
    print("Epoch: {} Accuracy: {}".format(epoch, acc))
sess.close()


